import numpy as np
from scipy.optimize import minimize, differential_evolution, NonlinearConstraint
import pandas as pd

# Read input data
df = pd.read_csv('sm.csv')

dfall = pd.read_csv('df.csv').drop(['G','P','tau','xbar'],axis=1)

dfall = dfall[['fbid','year','provname','name','deflator']]

df = pd.merge(df, dfall, on=['fbid','year'], how='left')

df['year_str'] = df['year'].astype(str)

df['SFE']=df['provname']+" "+df['name']+" "+df['year_str'] 

# Create a DataFrame to store sensitivity results
sensitivity_results = []

# Define parameter variation factors
factors = [0.8, 0.9, 1.0, 1.1, 1.2]
# Base parameters
b = 200
mu_base = 0.6
amen_base = 100
wbar = 10
f=4
S = 10e6
rr=1.1

checkbar=1000
checkbar1=100
xub=10000000
x1ub=500000


# Function to run optimization for given parameters
def run_optimization(G, p, tau, xbar, mu, amen, cp):
    c = cp * p

    def objective(vars):
        x, gamma = vars
        return -(wbar + b * (1 - 1 / (gamma * (1 - tau) * (p * x - c * x))) - f * (1 / (1 + xbar / (mu * x))))
    
    def constraint(vars):
        x, gamma = vars
        return 1 - G / ((1 - tau) * (p * x - c * x)) - gamma
     
    # Define bounds
    x_min = G / ((1 - tau) * (p - c)) + 100
    bounds = [(x_min, xub), (0, 1)]  # x in [x_min, 500000], gamma in [0, 1]
    
    # Define Nonlinear Constraint
    nl_constraint = NonlinearConstraint(constraint, 0, np.inf)
    
    # Step 1: Use Global Optimization for Initialization
    def global_optimization():
        result_global = differential_evolution(
            objective, 
            bounds=bounds, 
            constraints=(nl_constraint,),  # Pass the constraint properly
            seed=42
        )
        return result_global.x
    
    # Step 2: Local Optimization Refinement
    def local_optimization(initial_guess):
        options = {
            'maxiter': 1000,
            'ftol': 1e-6,
            'disp': True
        }
        result_local = minimize(
            objective,
            initial_guess,
            bounds=bounds,
            constraints={'type': 'ineq', 'fun': constraint},
            method='SLSQP',
            options=options
        )
        return result_local
    
    # Perform Global Optimization
    initial_guess = global_optimization()
    
    # Perform Local Optimization
    result = local_optimization(initial_guess)
    
    # Check Results
    if result.success:
        optimal_x, optimal_gamma = result.x
        optimal_value = -result.fun
        
        
    else:
        optimal_x=None
        optimal_gamma=None

    
    bbar=b*1.2
    
    def objective1(vars):
        x, gamma = vars
        return -(wbar + bbar * (1 - 1 / (gamma * (1 - tau) * (p * x - c * x))) + p * x - c * x+S-amen*x**rr)
    
    def constraint1(vars):
        x, gamma = vars
        return 1 - G / ((1 - tau) * (p * x - c * x)) - gamma
     
    # Define bounds
    x_min1 = G / ((1 - tau) * (p - c)) + 100
    bounds1 = [(x_min1, x1ub), (0, 1)]  # x in [x_min, 500000], gamma in [0, 1]
    
    # Define Nonlinear Constraint
    nl_constraint1 = NonlinearConstraint(constraint1, 0, np.inf)
    
    # Step 1: Use Global Optimization for Initialization
    def global_optimization1():
        result_global = differential_evolution(
            objective1, 
            bounds=bounds1, 
            constraints=(nl_constraint1,),  # Pass the constraint properly
            seed=42
        )
        return result_global.x
    
    # Step 2: Local Optimization Refinement
    def local_optimization1(initial_guess):
        options = {
            'maxiter': 1000,
            'ftol': 1e-6,
            'disp': True
        }
        result_local = minimize(
            objective1,
            initial_guess,
            bounds=bounds1,
            constraints={'type': 'ineq', 'fun': constraint1},
            method='SLSQP',
            options=options
        )
        return result_local
    
    # Perform Global Optimization
    initial_guess = global_optimization1()
    
    # Perform Local Optimization
    result1 = local_optimization1(initial_guess)
    
    def u(x):
        gamma=1 - G / ((1 - tau) * (p * x - c * x))    
        return wbar + b * (1 - 1 / (gamma * (1 - tau) * (p * x - c * x))) - f * (1 / (1 + xbar / (mu * x)))
        
    def check(x):
        
        if u(x-checkbar)<=optimal_value and u(x+checkbar)<=optimal_value:
            return 1   
        else:
            return 0
     
    def u1(x):
        gamma=1 - G / ((1 - tau) * (p * x - c * x))    
        return (wbar + bbar * (1 - 1 / (gamma * (1 - tau) * (p * x - c * x))) + p * x - c * x+S-amen*x**rr)
        
    def check1(x):
        if u1(x-checkbar1)<=optimal_value1 and u1(x+checkbar1)<=optimal_value1:
            return 1   
        else:
            return 0
            

    
 
    # Check Results
    if result1.success:
        optimal_x1, optimal_gamma1 = result1.x
        optimal_value1 = -result1.fun

        swl=bbar * (1 / ( optimal_gamma * (1 - tau) * (p *  optimal_x - c *  optimal_x))-1 / ( optimal_gamma1 * (1 - tau) * (p *  optimal_x1 - c *  optimal_x1))) + p *  optimal_x1 - c *  optimal_x1-(p *  optimal_x - c *  optimal_x)+amen*((optimal_x)**rr-(optimal_x1)**rr)
        
        if check(optimal_x) and check1(optimal_x1) and xub-optimal_x>1 and x1ub-optimal_x1>1 and optimal_x1<=xbar:
    
           return optimal_x, optimal_gamma, swl
       
        else:
            return None, None, None
    else:
        return None, None, None

n=len(df)
paralist=['mu', 'xbar', 'G','P','cp']
# Run sensitivity analysis for each parameter separately for all 58 observations
for factor in factors:
    for parameter in paralist:
        for i in range(n):
            try:
                print(i)
                G = df.loc[i, 'G'] if parameter != 'G' else df.loc[i, 'G'] * factor
                p = df.loc[i, 'P'] if parameter != 'P' else df.loc[i, 'P'] * factor
                cp=0.5411 if parameter != 'cp' else 0.5411 * factor
                tau = df.loc[i, 'tau']
                xbar = df.loc[i, 'xbar'] if parameter != 'xbar' else df.loc[i, 'xbar'] * factor
                mu = mu_base if parameter != 'mu' else mu_base * factor
                amen = amen_base if parameter != 'alpha' else amen_base * factor
                
                sfename=df.loc[i, 'SFE']
                deflator=df.loc[i, 'deflator']
    
                x, gamma, swl = run_optimization(G, p, tau, xbar, mu, amen,cp)
                sensitivity_results.append({
                    'observation': i,
                    'SFE': sfename,
                    'deflator': deflator,
                    'parameter': parameter,
                    'factor': factor,
                    'x': x,
                    'gamma': gamma,
                    'swl': swl
                })
            except ValueError as e:
                print(f"Skipping iteration {i} {parameter} {factor} due to error: {e}")
                continue
                

# Convert results to DataFrame
sensitivity_df = pd.DataFrame(sensitivity_results)

# Save results to a CSV file
sensitivity_df.to_csv('sensitivity_analysis_results_observations.csv', index=False)

